• dataclass.rst does not render methods.
• class.rst omits cythonized properties

cyclass places attributes section just after the main class docstring. this way we can document the actual attributes at the end of the main docstring and they are followed by docstring of all the properties.

We need to call decref in the except: branch.

The cleaners helpers do so. I moved the cpython.Py_INCREF call next to assigning its address to manager_ctx, for the clearer "release" ptr semantics. If the manager_ctx is set, the DECREF needs to be called.

TODO: cimport this from cydriver

🙃

I've actually had that this way initially, but seeing all the cdriver imports are gone from buffer, went along with the Python API. I can undig the previous variant.

Yeah, sorry. I think it's not "gone" gone, most likely @Andy-Jost found that we don't need many driver API calls in this file after the refactoring (#1205). But pointer attribute checks are in the hot path so we should cythonize it.

In fact, I am trying to catch up with what @fbusato is doing in C++ (NVIDIA/cccl#6733), which is an equivalent check (but for C++ mdspan instead of Python SMV).

Thanks for the reference! Looking at the logic in cccl, I adjusted managed memory "discovery". I am not sure if we need to go into so much details as trying to get particular memory pool and check if readability flag is set there, I did not add this, but can adjust if needed.

In any case, I moved back to cydriver API and added tests with host/device/managed/pinned from cuda malloc and pinned from cuda register.

For pinned memory, I am not 100% sure what happens on devices for which
CU_DEVICE_ATTRIBUTE_CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM is false. I.e. if one registers host memory with cuda host register and passes the original pointer, while the pointer to access on device is different. I.e. would it be still memory_type = 0 or memory_type = host and what would be a desired is_device_accessible value.

Q: is there any particular reason (apart from using 32-bit masks) we want to limit this to 32? There are libraries that use a higher limit, for example cuTENSOR (@yangcal would remember what the current limit is).

Initially, precompiling the copy kernels was on the table, so sticking to 32-ndim was more appealing. Now, I don't think there's anything preventing us from going to 64-dims. Especially, if there's really a need for that.

53568b5

@leofang,

NB.

from cuda.bindings import driver
driver.cuPointerGetAttributes(1, [driver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL], 0)

returns

(<CUresult.CUDA_SUCCESS: 0>, [4294967294])
while the actual value returned by the underlying API is -2. Looks like some unsigned/signed type mismatch.

How should we handle things like 6-bit and 4-bit floats? https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h#L171-L181

Good question, I am not sure, really. AFAIK, we don't support it now anyway, the dtype parsing won't accept it.

Playing with https://github.com/jax-ml/ml_dtypes I see that numpy:

1. does not allow exporting such types through dlpack
2. does not pack those - one element corresponds to one byte offset.

We could work around it in different ways, just looking at the landscape now I am not sure what's the desired behavior. E.g. would a dlpack dtype for fractional fps have nbits * lanes divisible by 8 or the numpy behavior is a common thing - we should round the itemsize to 1?

cc @seberg for visibility and any thoughts

Given JAX and ml_dtypes uses unpacked layouts and generally we want to support packed layouts that sounds like a bit of an impedance mismatch.

NumPy's problem with packed layout is that the array object's ABI is all byte-strided/pointer based. There might be paths to address this.¹
One could create e.g. a 2xfloat4_e2m1fn with explicit functions to pack and unpack. Not sure that a dtype helps a lot if you need to unpack for some or even most operations (compared to packing/unpacking with a custom converter function).

Enabling ml_dtypes to import/export for DLPack in NumPy, we can do, I think. Export would be easy, even a __dlpack_dtype__ would work. Unfortunately, a proper import implementation needs some form of registration so it probably not very quick.

FWIW, DLPack can allow bitsize == 8 to represent the padded versions. I don't think there was a particular reason not to just spell that out.
But by itself that doesn't actually quite help us (not without NumPy allowing export or some hacks here to make it work).

Thanks @seberg for detailed overview.

So when receiving fp4/fp6 tensor through dlpack we can expect 3 cases:

1. The "padded/unpacked" tensor (what numpy + ml_dtypes does now). We can represent them with no changes here, as soon as we start accepting the dtype. In dlpack, setting bitlen to 8 for those seems to be a bit of grey area, but if it's there, the StridedLayout won't be a hindrance. It looks like copying those does not require anything special.

2. The packed, vectorized dtype (2xfloat4_e2m1fn, 4xfloat6_e2m3fn?). We just need to take lanes into account, i.e. itemsize=nbits*lanes//8; assert itemsize * 8 == nbits*lanes. I omitted the lanes, because we forbid lanes > 1 when creating the numpy.dtype, but this is a simple adjustment. Again, it looks like copying those does not require anything special.

3. Receiving the packed, non-vectorized dtype.
   Little effort workaround would be to enforce scenario 2. I guess the actual tensors must have nice strides and round number of elements to make them performant, so this should be possible in any real-life case. We could either require users to pass views like in scenario 2 or stipulate that we do such a conversion and repack the layout automatically. If really necessary we could special-case the SMV to keep the original shape and strides and use the repacked StridedLayout only internally if a copy is requested.

These checks seem somewhat expensive to do on the path of getting the ptr to the data. Should we validate these only at construction time and avoid them in a hot path like this?

Those checks happen now for new SMV created with from_buffer or view methods. We could limit those to from_buffer only or get rid of them completely. I don't have strong preference either way here.

My 2c:

• If Buffer isn't validating ownership, neither should SMV. If we wanted to require an owner we should do so at the Buffer level. Let's remove the is_allocated check entirely?
• Given we're allowing effectively slicing a buffer at SMV construction time using offset and size, I think the layout size validation makes sense at from_buffer and view construction time.

Can we call cuInit if that's what we're after, instead of calling Device?

We're already importing numpy where this is being used, can't we use np.s_ for this instead of duplicating that functionality solely for testing?

Cool, I wasn't aware of this utility.